# Worldviews
# ..............................................................................

import matplotlib.pyplot as plt
import networkx as nx
from pyvis import network as net


def plot_interactive_network(
    networkx_graph,
    notebook=False,
    output_filename="graph.html",
    show_buttons=False,
    only_physics_buttons=False,
    gravity = -100,
    options = None,
    width="1000px",
    height="1000px",
):
    
    # make a pyvis network
    pyvis_graph = net.Network(notebook=notebook, directed=True)
    pyvis_graph.width = width
    pyvis_graph.height = height

    # for each node and its attributes in the networkx graph
    for node, node_attrs in networkx_graph.nodes(data=True):
        pyvis_graph.add_node(node, **node_attrs)

    # for each edge and its attributes in the networkx graph
    for source, target, edge_attrs in networkx_graph.edges(data=True):
        pyvis_graph.add_edge(source, target, **edge_attrs)

    # turn buttons on
    if show_buttons:
        if only_physics_buttons:
            pyvis_graph.show_buttons(filter_=["physics"])
        else:
            pyvis_graph.show_buttons()

    # Make sure edges aren't written on one another
    pyvis_graph.set_edge_smooth("dynamic")

    # Repulsion settings
    pyvis_graph.barnes_hut(gravity=gravity, overlap=0.1)
    
    # Options
    if options:
        pyvis_graph.set_options(
            options
        )
    
    # return and also save
    return pyvis_graph.show(output_filename)


import pandas as pd
from datetime import datetime

c = 1000

df = pd.read_csv('../data/gpo_final_data/narratives_complete_with_metadata_manual_labels_rich_{0}.csv'.format(c))

# Represent the top K partisan narratives in a directed multigraph
# ..............................................................................

K = 100
prune_network=True

temp = df[(df.party == 'Republican') | (df.party == 'Democrat')]

temp = temp[['ARGO', 'B-V-RAW', 'ARG1', 'narrative', 'log_or', 'log_or_ci_lower', 'log_or_ci_upper', 'frequency']].drop_duplicates()
temp['B-V'] = temp['B-V-RAW']

temp = temp.sort_values(by='frequency', ascending=False).iloc[0:K]
temp['q_freq'] = pd.qcut(temp['frequency'], 30, duplicates='drop', labels=False)+1

temp = temp.to_dict(orient = 'records')

list_of_args = set.union(set(df['ARGO']), set(df['ARG1']))

G = nx.MultiDiGraph()

for l in temp:    
    if l['log_or_ci_lower'] <= 0 <= l['log_or_ci_upper']:
        l['color'] = "rgba(131, 133, 143, 0.7)" #'grey'
    else:
        if l['log_or'] >= 0:
            l['color'] = "rgba(187, 44, 44, 0.7)" #'red'
        else:
            l['color'] = "rgba(36, 75, 153, 0.7)"#'blue'
            
    G.add_edge(
        l["ARGO"],
        l["ARG1"],
        width=l['q_freq']/3,
        label=l["B-V"],
        hidden=False,
        color=l["color"]
    )
        
d = nx.degree(G)  
degrees = [G.degree(n) for n in G.nodes()]
degrees = pd.qcut(degrees, 10, duplicates='drop', labels=False)+1

for i,node in enumerate(list(G.nodes)):
    G.nodes[node]["size"] = d[node]
    G.nodes[node]["color"] = "rgb(166, 187, 190, 0.7)"
    G.nodes[node]["hidden"] = False

if prune_network:
    # Generate connected components and select the largest:
    largest_component = max(nx.weakly_connected_components(G), key=len)
    # Create a subgraph of G consisting only of this component:
    G = G.subgraph(largest_component)
    
options =     """
    var options = {
      "nodes": {
        "font": {
          "size": 20
        }
      },
      "edges": {
        "color": {
          "inherit": true
        },
        "font": {
          "size": 20
        },
        "smooth": {
          "forceDirection": "none"
        }
      },
      "physics": {
        "barnesHut": {
          "centralGravity": 0.1,
          "springLength": 100,
          "springConstant": 0.1,
          "avoidOverlap": 0.4
        },
        "minVelocity": 0.75
      }
    }
    """

plot_interactive_network(
    G,
    notebook = False,
    output_filename = '../figures/Figure_4.html',
    gravity = -1000,
    options = options,
    width="1000px",
    height="1000px",
    )

options =     """
    var options = {
      "nodes": {
        "font": {
          "size": 20
        }
      },
      "edges": {
        "color": {
          "inherit": true
        },
        "font": {
          "size": 20
        },
        "smooth": {
          "forceDirection": "none"
        }
      },
      "physics": {
        "barnesHut": {
          "centralGravity": 0.5,
          "springLength": 100,
          "springConstant": 0.05,
          "avoidOverlap": 0.1
        },
        "minVelocity": 0.75
      }
    }
    """

# Republican narratives

K = 50
temp = df[(df.party == 'Republican') | (df.party == 'Democrat')]
temp = temp[['ARGO', 'B-V-RAW', 'ARG1', 'narrative', 'log_or', 'log_or_ci_lower', 'log_or_ci_upper', 'frequency']].drop_duplicates()
temp['B-V'] = temp['B-V-RAW']

temp = temp.sort_values(by='log_or', ascending=False).iloc[0:K]
temp['q_freq'] = pd.qcut(temp['log_or'], 30, duplicates='drop', labels=False)+1

temp = temp.to_dict(orient = 'records')

list_of_args = set.union(set(df['ARGO']), set(df['ARG1']))

G = nx.MultiDiGraph()

for l in temp:    
    l['color'] = "rgba(187, 44, 44, 0.7)" #red
            
    G.add_edge(
        l["ARGO"],
        l["ARG1"],
        width=l['q_freq']/3,
        label=l["B-V"],
        hidden=False,
        color=l["color"]
    )
        
d = nx.degree(G)  
degrees = [G.degree(n) for n in G.nodes()]
degrees = pd.qcut(degrees, 10, duplicates='drop', labels=False)+1

for i,node in enumerate(list(G.nodes)):
    G.nodes[node]["size"] = d[node]*3
    G.nodes[node]["color"] = "rgb(166, 187, 190, 0.7)"
    G.nodes[node]["hidden"] = False

if prune_network:
    # Generate connected components and select the largest:
    largest_component = max(nx.weakly_connected_components(G), key=len)
    # Create a subgraph of G consisting only of this component:
    G = G.subgraph(largest_component)
    
plot_interactive_network(
    G,
    notebook = False,
    output_filename = '../figures/Figure_H_3.html',
    gravity = -100,
    options = options,
    width="1000px",
    height="1000px",
    )

# Democrat narratives

K = 50
temp = df[(df.party == 'Republican') | (df.party == 'Democrat')]
temp = temp[['ARGO', 'B-V-RAW', 'ARG1', 'narrative', 'log_or', 'log_or_ci_lower', 'log_or_ci_upper', 'frequency']].drop_duplicates()
temp['B-V'] = temp['B-V-RAW']

temp = temp.sort_values(by='log_or', ascending=True).iloc[0:K]
temp['q_freq'] = pd.qcut(temp['log_or'], 30, duplicates='drop', labels=False)+1

temp = temp.to_dict(orient = 'records')

list_of_args = set.union(set(df['ARGO']), set(df['ARG1']))

G = nx.MultiDiGraph()

for l in temp:    
    l['color'] = "rgba(36, 75, 153, 0.7)" #blue
            
    G.add_edge(
        l["ARGO"],
        l["ARG1"],
        width=l['q_freq']/3,
        label=l["B-V"],
        hidden=False,
        color=l["color"]
    )
        
d = nx.degree(G)  
degrees = [G.degree(n) for n in G.nodes()]
degrees = pd.qcut(degrees, 10, duplicates='drop', labels=False)+1

for i,node in enumerate(list(G.nodes)):
    G.nodes[node]["size"] = d[node]*2
    G.nodes[node]["color"] = "rgb(166, 187, 190, 0.7)"
    G.nodes[node]["hidden"] = False

if prune_network:
    # Generate connected components and select the largest:
    largest_component = max(nx.weakly_connected_components(G), key=len)
    # Create a subgraph of G consisting only of this component:
    G = G.subgraph(largest_component)
    
plot_interactive_network(
    G,
    notebook = False,
    output_filename = '../figures/Figure_H_2.html',
    gravity = -50,
    options = options,
    width="1000px",
    height="1000px",
    )
